package com.cyber4aiagent.config;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import com.cyber4aiagent.store.MongoChatMemoryStore;

import java.util.Arrays;
import java.util.List;

/**
 * @author Cyber
 * @date 2025/05/09 18:09
 * @description 小智智能体配置
 */
@Configuration
public class SmallZhiAgentAssistantConfig {

    @Autowired
    private MongoChatMemoryStore mongoChatMemoryStore;

    /**
     * 隔离聊天记忆功能的提供者
     */
    @Bean
    public ChatMemoryProvider chatMemorySmallZhiProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(30)
                .chatMemoryStore(mongoChatMemoryStore)
                .build();
    }

    /**
     * 内容检索器
     */
    @Bean
    public ContentRetriever contentRetrieverSmallZhi() {
        // 使用FileSystemDocumentLoader读取指定目录下的知识库文档
        // 并使用默认的文档解析器TextDocumentParser对文档进行解析
        Document document = FileSystemDocumentLoader.loadDocument("src/main/resources/small-zhi/test.md");
        Document document1 = FileSystemDocumentLoader.loadDocument("src/main/resources/small-zhi/test1.md");
        Document document2 = FileSystemDocumentLoader.loadDocument("src/main/resources/small-zhi/test2.md");

        List<Document> documents = Arrays.asList(document, document1, document2);

        // 使用内存向量数据库存储知识库文档
        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

        // 使用默认的文档分割器
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);

        // 从嵌入存储 (EmbeddingStore) 里检索和查询内容相关的信息
        return EmbeddingStoreContentRetriever.from(embeddingStore);
    }

    /**
     * 添加基于 pinecone 的向量数据库配置
     */
    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private EmbeddingStore embeddingStore;

    @Bean
    public ContentRetriever contentRetrieverSmallZhiPineCone() {
        // 从嵌入存储 (EmbeddingStore) 里检索和查询内容相关的信息
        return EmbeddingStoreContentRetriever.builder()
                // 设置用于生成嵌入向量的嵌入模型
                .embeddingModel(embeddingModel)
                // 指定要使用的嵌入存储
                .embeddingStore(embeddingStore)
                // 设置最大的检索结果数量，这里表示最多返回1条匹配结果
                .maxResults(1)
                // 设置最小的匹配分数，这里表示匹配分数必须大于等于0.8，才会返回
                .minScore(0.8)
                .build();
    }
}
